package com.witmore.neutron.ai.cognitive.advisor;

import com.witmore.neutron.ai.cognitive.dto.SessionDeductionRequest;
import com.witmore.neutron.ai.cognitive.service.UserSessionUsageService;
import com.witmore.neutron.ai.db.entity.ModelConfigDO;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.ai.chat.client.ChatClientRequest;
import org.springframework.ai.chat.client.ChatClientResponse;
import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
import org.springframework.ai.chat.client.advisor.api.StreamAdvisor;
import org.springframework.ai.chat.client.advisor.api.StreamAdvisorChain;
import org.springframework.ai.chat.metadata.Usage;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.stereotype.Service;
import org.springframework.util.ObjectUtils;
import reactor.core.publisher.Flux;

import java.util.Map;
import java.util.Objects;
import java.util.concurrent.atomic.AtomicReference;

/**
 * Description: CustomSimpleLoggerAdvisor
 * Author: 豆子高
 * Date: 2025/9/5 17:26
 */
@Slf4j
@Service
public class CustomSessionUsageAdvisor implements CallAdvisor, StreamAdvisor {


    private final UserSessionUsageService userSessionUsageService;

    public CustomSessionUsageAdvisor(UserSessionUsageService userSessionUsageService) {
        this.userSessionUsageService = userSessionUsageService;
    }

    /**
     * 获取会话扣减请求参数
     * & 校验会话次数
     *
     * @param context 上下文信息
     * @return 扣减请求参数
     */
    private SessionDeductionRequest getSessionRequest(Map<String, Object> context) throws Exception {
        // 获取上下文信息
        Integer userId = (Integer) context.get("userId");
        ModelConfigDO modelConfigDO = (ModelConfigDO) context.get("modelConfigDO");
        String sessionId = (String) context.get("sessionId");
        Integer agentId = (Integer) context.get("agentId");
        ToolCallback[] toolCallbacks = (ToolCallback[]) context.get("toolCallbacks");
        String userApiKey = (String) context.get("userApiKey");
        log.info("获取上下文信息 - context: {}", ObjectUtils.isEmpty(context) ? "空" : context);
        if (userId == null || modelConfigDO == null) {
            log.warn("缺少必要的上下文信息 - 用户ID: {}, 模型信息: {}", userId, modelConfigDO);
            throw new Exception("系统异常：缺少必要的上下文信息");
        }
        // 构建扣减请求参数
        SessionDeductionRequest request = SessionDeductionRequest.of(
                userId, modelConfigDO, sessionId, agentId, toolCallbacks, userApiKey);
        // 先校验次数
        boolean checked = userSessionUsageService.checkUserSessionCount(request);
        if (!checked) {
            throw new Exception("会话次数不足，请充值后再试");
        }
        return request;
    }

    @NotNull
    @Override
    public ChatClientResponse adviseCall(@NotNull ChatClientRequest request, @NotNull CallAdvisorChain chain) {
        log.info("CustomSessionUsageAdvisor 次数记录 请求开始: {}", request);
        try {
            // 获取上下文信息
            SessionDeductionRequest deductionRequest = getSessionRequest(request.context());
            // 执行对话请求
            ChatClientResponse response = chain.nextCall(request);
            Usage usage = Objects.requireNonNull(response.chatResponse()).getMetadata().getUsage();
            deductionRequest.setUsageTokensIn(usage.getPromptTokens());
            deductionRequest.setUsageTokensOut(usage.getCompletionTokens());
            deductionRequest.setUsageTotalTokens(usage.getTotalTokens());
            String rowMetadataId = response.chatResponse().getMetadata().getId();
            log.info("用户 {} 会话次数记录成功，rowMetadataID: {}", deductionRequest.getUserId(), rowMetadataId);
            deductionRequest.setRowMetadataId(rowMetadataId);
            // 直接进行扣减操作，内部会原子地完成检查+扣减，并记录完整信息
            userSessionUsageService.deductUserSessionCount(deductionRequest);
            log.info("用户 {} 会话次数扣减成功", deductionRequest.getUserId());
            log.info("请求完成，扣减成功！用户: {}", deductionRequest.getUserId());
            return response;
        } catch (Exception e) {
            throw new RuntimeException(e.getMessage());
        }
    }

    @NotNull
    @Override
    public Flux<ChatClientResponse> adviseStream(@NotNull ChatClientRequest request, @NotNull StreamAdvisorChain chain) {
        log.info("CustomSessionUsageAdvisor 次数记录 流式请求开始: {}", request);

        try {
            // 获取上下文信息
            SessionDeductionRequest deductionRequest = getSessionRequest(request.context());
            // 执行流式请求
            Flux<ChatClientResponse> chatClientResponseFlux = chain.nextStream(request);
            AtomicReference<ChatClientResponse> lastResponse = new AtomicReference<>();
            // 对流式响应进行处理和日志记录
            return chatClientResponseFlux
                    .doOnNext(response -> {
                        lastResponse.set(response);
                        log.debug("收到流式响应片段，用户: {}", deductionRequest.getUserId());
                    })
                    .doOnComplete(() -> {
                        // 流结束时，检查最后一个响应块中是否有usage信息
                        ChatClientResponse finalResponse = lastResponse.get();
                        if (finalResponse != null) {
                            Usage finalUsage = Objects.requireNonNull(finalResponse.chatResponse()).getMetadata().getUsage();
                            if (finalUsage != null) {
                                log.info("Final Usage Stream: {}", finalUsage);
                                deductionRequest.setUsageTokensIn(finalUsage.getPromptTokens());
                                deductionRequest.setUsageTokensOut(finalUsage.getCompletionTokens());
                                deductionRequest.setUsageTotalTokens(finalUsage.getTotalTokens());
                            }
                            String rowMetadataId = finalResponse.chatResponse().getMetadata().getId();
                            log.info("MetadataId: {}", rowMetadataId);
                            deductionRequest.setRowMetadataId(rowMetadataId);
                        }
                        // 直接进行扣减操作，内部会原子地完成检查+扣减，并记录完整信息
                        userSessionUsageService.deductUserSessionCount(deductionRequest);
                        log.info("用户 {} 会话次数扣减成功", deductionRequest.getUserId());
                        log.info("流式请求完成，扣减成功！用户: {}", deductionRequest.getUserId());
                    })
                    .doOnError(error -> log.error("流式请求异常，用户: {}", deductionRequest.getUserId(), error));
        } catch (Exception e) {
            return Flux.error(new RuntimeException(e.getMessage()));
        }

    }

    @NotNull
    @Override
    public String getName() {
        return this.getClass().getSimpleName();
    }

    /**
     * 2025-09-10
     * 目前有两个advisor 优先级order数值越低 before越先执行 after越后执行
     * 1. 自定义的CustomSessionUsageAdvisor 优先级高一些 用于校验调用可行性 主要校验次数是否够用 同时记录会话次数的
     * 2. 系统自带的PromptChatMemoryAdvisor  优先级低 用于持久化聊天记录的
     */
    @Override
    public int getOrder() {
        return 0;
    }
}
